from os.path import join
from pymarl.src.envs import REGISTRY as env_REGISTRY
from functools import partial
from pymarl.src.components.episode_buffer import EpisodeBatch

import numpy as np
import pickle
import time
import os


class EpisodeRunner:
    def __init__(self, args, logger):
        self.args = args
        self.logger = logger
        self.batch_size = self.args.batch_size_run
        assert self.batch_size == 1

        self.env = env_REGISTRY[self.args.env](**self.args.env_args)
        self.episode_limit = self.env.episode_limit
        self.t = 0
        self.t_env = 0

        self.train_returns = []
        self.test_returns = []
        self.train_stats = {}
        self.test_stats = {}

        self.log_train_stats_t = -1000000

        os.makedirs("offline_data", exist_ok=True)

    def setup(self, scheme, groups, preprocess, mac):
        self.new_batch = partial(EpisodeBatch, scheme, groups, self.batch_size, self.episode_limit + 1,
                                 preprocess=preprocess, device=self.args.device)
        self.mac = mac

    def get_env_info(self):
        return self.env.get_env_info()

    def save_replay(self):
        self.env.save_replay()

    def close_env(self):
        self.env.close()

    def reset(self):
        self.batch = self.new_batch()
        self.env.reset()
        self.t = 0

    def run(self, test_mode=False):
        self.reset()

        terminated = False
        episode_return = 0
        self.mac.init_hidden(batch_size=self.batch_size)

        episode_trajectory = [[] for _ in range(self.env.n_agents)]

        while not terminated:
            pre_transition_data = {
                "state": [self.env.get_state()],
                "avail_actions": [self.env.get_avail_actions()],
                "obs": [self.env.get_obs()]
            }

            self.batch.update(pre_transition_data, ts=self.t)
            actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
            reward, terminated, env_info = self.env.step(actions[0])
            episode_return += reward

            state = self.env.get_state()
            obs = self.env.get_obs()
            avail_actions = self.env.get_avail_actions()

            # 正确获取 kill_count 和 is_alive
            # === 提取 kill_count 和 is_alive（从 info 中取） ===
            kill_counts = env_info.get("kill_count", [0] * self.env.n_agents)
            is_alives = env_info.get("is_alive", [1] * self.env.n_agents)

            # === 记录每个 agent 的 step 数据 ===
            for agent_id in range(self.env.n_agents):
                step = [
                    state,  # global state
                    obs[agent_id],  # agent-specific obs
                    [actions[0][agent_id].item()],  # action
                    [reward],  # reward
                    terminated,  # done
                    avail_actions[agent_id],  # avail actions
                    [kill_counts[agent_id]],  # 使用环境 info 中的击杀数（累计）
                    [is_alives[agent_id]]  # 存活状态（1 或 0）
                ]
                episode_trajectory[agent_id].append(step)

            post_transition_data = {
                "actions": actions,
                "reward": [(reward,)],
                "terminated": [(terminated != env_info.get("episode_limit", False),)],
            }

            self.batch.update(post_transition_data, ts=self.t)
            self.t += 1

        last_data = {
            "state": [self.env.get_state()],
            "avail_actions": [self.env.get_avail_actions()],
            "obs": [self.env.get_obs()]
        }
        self.batch.update(last_data, ts=self.t)
        actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
        self.batch.update({"actions": actions}, ts=self.t)

        # ✅ 保存为 .pkl（供后续 json 转换）
        save_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "offline_data"
        )
        os.makedirs(save_dir, exist_ok=True)

        save_path = os.path.join(save_dir, f"episode_{time.time()}.pkl")
        with open(save_path, "wb") as f:
            pickle.dump(episode_trajectory, f)

        print(f"[EPISODE] PKL轨迹已保存到: {save_path}")

        # Logging
        cur_stats = self.test_stats if test_mode else self.train_stats
        cur_returns = self.test_returns if test_mode else self.train_returns
        log_prefix = "test_" if test_mode else ""

        for k in set(cur_stats) | set(env_info):
            env_value = env_info.get(k)
            if isinstance(env_value, list):
                if k not in cur_stats:
                    cur_stats[k] = env_value.copy()
                else:
                    min_len = min(len(cur_stats[k]), len(env_value))
                    cur_stats[k] = [cur_stats[k][i] + env_value[i] for i in range(min_len)]
                    if len(cur_stats[k]) < len(env_value):
                        cur_stats[k].extend(env_value[min_len:])
            else:
                cur_stats[k] = cur_stats.get(k, 0) + (env_value if env_value is not None else 0)

        cur_stats["n_episodes"] = 1 + cur_stats.get("n_episodes", 0)
        cur_stats["ep_length"] = self.t + cur_stats.get("ep_length", 0)

        if not test_mode:
            self.t_env += self.t

        cur_returns.append(episode_return)

        if test_mode and (len(self.test_returns) == self.args.test_nepisode):
            self._log(cur_returns, cur_stats, log_prefix)
        elif self.t_env - self.log_train_stats_t >= self.args.runner_log_interval:
            self._log(cur_returns, cur_stats, log_prefix)
            if hasattr(self.mac.action_selector, "epsilon"):
                self.logger.log_stat("epsilon", self.mac.action_selector.epsilon, self.t_env)
            self.log_train_stats_t = self.t_env

        return self.batch

    def _log(self, returns, stats, prefix):
        self.logger.log_stat(prefix + "return_mean", np.mean(returns), self.t_env)
        self.logger.log_stat(prefix + "return_std", np.std(returns), self.t_env)
        returns.clear()

        n_episodes = stats["n_episodes"]
        for k, v in stats.items():
            if k == "n_episodes":
                continue

            if isinstance(v, list):
                for i, value in enumerate(v):
                    self.logger.log_stat(
                        f"{prefix}{k}_agent{i}_mean",
                        value / n_episodes,
                        self.t_env
                    )
            else:
                self.logger.log_stat(
                    f"{prefix}{k}_mean",
                    v / n_episodes,
                    self.t_env
                )

        stats.clear()
